#!/usr/bin/env python
# -*- coding:utf-8 -*-
'''
@File   :   dataloader.py
@Author :   Song
@Time   :   2022/2/28 21:47
@Contact:   songjian@westlake.edu.cn
@intro  : 
'''
import torch
import pandas as pd
import numpy as np
import pandas as pd
from torch.utils.data.dataset import Dataset
import predifine
import operator
from pathlib import Path

class Train_Dataset(Dataset):

    def __init__(self, info, path):
        # load
        npz = np.load(path)
        self.pr_mz = npz['pr_mz']
        self.pr_charge = npz['pr_charge']
        self.scan_x = npz['scan_x']
        self.scan_y = npz['scan_y']
        scan_len = npz['scan_len']

        # 统计正负样本谱图
        self.label = npz['label']
        print(f'{info}, pos spectra num: {sum(self.label == 1)}, neg spectra num: {sum(self.label == 0)}')

        # 处理scan信号的索引
        self.scan_access_idx = [0]
        self.scan_access_idx.extend(scan_len)
        self.scan_access_idx = np.cumsum(self.scan_access_idx)

    def __len__(self):
        return len(self.pr_mz)
        # return 1000000

    def __getitem__(self, idx):
        idx = np.random.randint(len(self.pr_mz))

        # 处理scan
        begin, end = self.scan_access_idx[idx], self.scan_access_idx[idx + 1]
        scan_x = self.scan_x[begin: end]
        scan_y = self.scan_y[begin: end]
        scan_y = scan_y / scan_y.max()

        # 处理pr
        pr_mz = self.pr_mz[idx]
        pr_charge = self.pr_charge[idx]

        # 把pr嵌入到scan，强度补1.
        scan_x = np.append(scan_x, pr_mz)
        scan_y = np.append(scan_y, 1.)
        spectra = np.column_stack([scan_x, scan_y])

        # A数目
        label = self.label[idx]

        return (torch.tensor(spectra, dtype=torch.float),
                pr_charge,
                label)



class Eval_Dataset(Dataset):

    def __init__(self, info, path):
        # load
        npz = np.load(path)
        self.pr_mz = npz['pr_mz']
        self.pr_charge = npz['pr_charge']
        self.scan_x = npz['scan_x']
        self.scan_y = npz['scan_y']
        scan_len = npz['scan_len']

        # 统计正负样本谱图
        self.label = npz['label']
        print(f'{info}, pos spectra num: {sum(self.label == 1)}, neg spectra num: {sum(self.label == 0)}')

        # 处理scan信号的索引
        self.scan_access_idx = [0]
        self.scan_access_idx.extend(scan_len)
        self.scan_access_idx = np.cumsum(self.scan_access_idx)

    def __len__(self):
        return len(self.pr_mz)

    def __getitem__(self, idx):
        # 处理scan
        begin, end = self.scan_access_idx[idx], self.scan_access_idx[idx + 1]
        scan_x = self.scan_x[begin: end]
        scan_y = self.scan_y[begin: end]
        scan_y = scan_y / scan_y.max()

        # 处理pr
        pr_mz = self.pr_mz[idx]
        pr_charge = self.pr_charge[idx]

        # 把pr嵌入到scan，强度补1.
        scan_x = np.append(scan_x, pr_mz)
        scan_y = np.append(scan_y, 1.)
        spectra = np.column_stack([scan_x, scan_y])

        # A数目
        label = self.label[idx]

        return (torch.tensor(spectra, dtype=torch.float),
                pr_charge,
                label)


class mzML_Dataset(Dataset):

    def __init__(self, mzml, test_idx):
        # load
        self.pr_mz = mzml.all_pr_mz[test_idx]
        self.pr_charge = mzml.all_pr_charge[test_idx]
        self.scan_x = np.concatenate(mzml.all_mz[test_idx])
        self.scan_y = np.concatenate(mzml.all_intensity[test_idx])
        scan_len = mzml.all_scan_len[test_idx]

        assert sum(scan_len) == len(self.scan_x) == len(self.scan_y)
        print(f'mzML原有二级谱图{sum(mzml.all_levels == 2)}张，需要打分{len(self.pr_mz)}张。')

        # 处理scan信号的索引
        self.scan_access_idx = [0]
        self.scan_access_idx.extend(scan_len)
        self.scan_access_idx = np.cumsum(self.scan_access_idx)

    def __len__(self):
        return len(self.pr_mz)

    def __getitem__(self, idx):
        # 处理scan
        begin, end = self.scan_access_idx[idx], self.scan_access_idx[idx + 1]
        scan_x = self.scan_x[begin: end]
        scan_y = self.scan_y[begin: end]
        scan_y = scan_y / scan_y.max()

        # 处理pr
        pr_mz = self.pr_mz[idx]
        pr_charge = self.pr_charge[idx]

        # 把pr嵌入到scan，强度补1.
        scan_x = np.append(scan_x, pr_mz)
        scan_y = np.append(scan_y, 1.)
        spectra = np.column_stack([scan_x, scan_y])

        # A数目
        label = -1

        return (torch.tensor(spectra, dtype=torch.float),
                pr_charge,
                label)
